AIMET Integration by kozlov721 · Pull Request #366 · luxonis/luxonis-train

kozlov721 · 2026-04-08T07:24:10Z

Purpose

Adds option to quantize trained model using various PTQ and QAT techniques.

Specification

Added LuxonisModel.quantize method
- Runs selected algorithms on the model weights and exports it to ONNX
- Support for various PQT techniques:
  - Adaround
  - Batch norm folding
  - Cross-layer equalization
  - Batch norm re-estimation
  - Sequential MSE
Added AIMETCallback
- Runs quantization at the end of the training
Added a new section exporter.aimet to the Config
Other small changes necessary to accomodate the AIMET API or to simplify the integration
- Added more standard forward method to LuxonixLightningModule
  - The original renamed to full_forward
- Support for loaders outputting bare tensors
  - As opposed to the expected dictionary
- Removed metrics, losses and visuazilers from saved node modules
- Added print_table as a required abstract method to BaseLuxonisProgressBar
- Implemented custom __getstate__ and __setstate__ in LuxonisLightningModule

config.yaml

# Default values
exporter:
  aimet:
    active: false

    default_output_bw: 8
    default_param_bw: 8
    default_data_type: "int"
    quant_scheme: "min_max"
    config: ~

    fold_batch_norms: false
    cross_layer_equalization: false
    batch_norm_reestimation: false
    sequential_mse: false
    adaround:
      active: false
      default_num_iterations: ~
      default_reg_param: 0.01
      default_beta_range: [20, 2]
      default_warm_start: 0.2

    epochs: 20
    optimizer:
      name: "SGD"
      params:
        lr: 0.00001
    scheduler:
      name: "StepLR"
      params:
        step_size: 5
        gamma: 0.1

Dependencies & Potential Impact

None / not applicable

Deployment Plan

None / not applicable

Testing & Validation

Testing AIMETCallback in test_callbacks
Testing LuxonisModel.quantize for all predefined models with the full set of PTQ techniques enabled

Summary by CodeRabbit

New Features
- AIMET quantization: new quantize CLI command and model.quantize API; optional AIMET install instructions included.
Bug Fixes & Improvements
- Loaders now accept raw tensors or image-keyed dicts.
- Improved export/pickling, progress-bar table rendering, visualization stability, and training callback behavior.
- Model initialization accepts weights from the CLI.
Documentation
- Installation and CLI docs updated to document AIMET/quantize.

kozlov721

could you also please include in the PR comment a few notes/justifications on the changes that are not necessarily related to the PR?

Mostly all the changes are somehow related to making AIMET work. I added more comments to some individual changes in code. If you have more questions about some other specific changes let me know.

coderabbitai · 2026-05-21T10:29:14Z

Warning

Review limit reached

@kozlov721, we couldn't start this review because you've reached your PR review rate limit.

More reviews will be available in 6 minutes and 29 seconds. Learn how PR review limits work.

Your organization has run out of usage credits. Purchase more in the billing tab.

⌛ How to resolve this issue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans include higher PR review limits than trial, open-source, and free plans. In all cases, reviews become available again over time. During sustained high-volume PR review activity, CodeRabbit may temporarily slow when the next review becomes available.

Please see our Fair Usage Limits Policy for further information.

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 7a609f2b-d30a-4bef-9da5-76841d8dc558

📥 Commits

Reviewing files that changed from the base of the PR and between 8d25cf2 and 42370f1.

📒 Files selected for processing (15)

.github/workflows/ci.yaml
README.md
configs/README.md
luxonis_train/__main__.py
luxonis_train/callbacks/README.md
luxonis_train/config/config.py
luxonis_train/core/core.py
luxonis_train/loaders/base_loader.py
luxonis_train/loaders/luxonis_loader_torch.py
requirements.txt
tests/integration/test_cli_commands.py
tests/integration/test_combinations.py
tests/integration/test_core.py
tests/integration/test_custom_model.py
tests/integration/test_fixed_validation_batch_limit.py

📝 Walkthrough

Walkthrough

Adds AIMET quantization support (config, utilities, LuxonisModel.quantize, CLI, callback) and refactors loader/forward interfaces to accept dict or tensor inputs; updates Lightning export/pickling, loss/visualizer/metric modules, packaging, CI, docs, and tests.

Changes

AIMET Quantization Integration

Layer / File(s)	Summary
AIMET config and models `luxonis_train/config/config.py`, `pyproject.toml`, `requirements-aimet.txt`	Adds `AdaroundConfig` and `AIMETConfig`, enum handling, `ExportConfig.aimet`, and `aimet` optional extra with requirements file.
AIMET utilities `luxonis_train/core/utils/aimet_utils.py`	Adds `check_aimet_available`, `post_training_quantization`, and `quantization_aware_training` implementing PTQ/QAT flows and AdaRound handling.
LuxonisModel.quantize & core wiring `luxonis_train/core/core.py`	Implements `LuxonisModel.quantize()` end-to-end, weights CLI handling, typed loader properties, and test/thread return behavior fixes.
CLI command and AIMETCallback `luxonis_train/__main__.py`, `luxonis_train/callbacks/aimet_callback.py`, `luxonis_train/callbacks/__init__.py`	Adds `quantize` CLI subcommand, registers `AIMETCallback`, and exports it for use when enabled.
CI, packaging, docs `.github/workflows/ci.yaml`, `README.md`, `configs/README.md`, `luxonis_train/callbacks/README.md`	CI installs `[dev,aimet]` with PyTorch cu130 index; README and configs document AIMET extras, quantize CLI, and exporter AIMET options.
Loader API refactor `luxonis_train/loaders/base_loader.py`, `luxonis_train/loaders/`, `tests/`, `media/*`	Switches from `get()` to `__getitem__`, widens loader output to `dict[str,Tensor]
Input-normalization & inference paths `luxonis_train/core/utils/infer_utils.py`, `luxonis_train/core/utils/annotate_utils.py`, `luxonis_train/__main__.py`, `luxonis_train/callbacks/gradcam_visualizer.py`	Use `full_forward`, normalize raw Tensor inputs to `{image_source: tensor}`, and handle non-dict images in visualization extraction.
Lightning forward/export/pickling `luxonis_train/lightning/luxonis_lightning.py`	Changes `forward` to accept dict
Losses: buffer registration & order fixes `luxonis_train/attached_modules/losses/*`	Registers `gt_bboxes_scale`/`gt_kpts_scale` as non-persistent buffers in init, reorders init usage, and uses cloned stride tensors.
Metrics & Visualizers `luxonis_train/attached_modules/metrics/`, `luxonis_train/attached_modules/visualizers/`	Adds identity-based equality/hash for BaseMetric, clones confusion matrix when inference-mode, centralizes `colormap` in BaseVisualizer with `__getstate__`, sanitizes embeddings (NaN/Inf), and removes local ColorMap instantiation.
Progress bar table refactor `luxonis_train/callbacks/luxonis_progress_bar.py`	Introduces abstract `print_table` and updates TQDM/Rich implementations to use iterable rows and column names; formats floats for Rich output.
Quantization-compatible blocks & hub flags `luxonis_train/nodes/backbones/pplcnet_v3/blocks.py`, `luxonis_train/nodes/blocks/__init__.py`, `luxonis_train/nodes/backbones/*`	Adds conditional `QuantizedAffineBlock` under AIMET guard, configures mixin ignores, and sets `trust_repo=True` for torch.hub loads.
Misc improvements `luxonis_train/nodes/blocks/blocks.py`, `luxonis_train/registry.py`, `luxonis_train/utils/`, `luxonis_train/config/predefined_models/`	Make reparametrize/restore idempotent, add `_INTERNAL` registry, dedupe `__all__`, refine main_metric auto-selection, and improve dataset metadata messages.
Tests updated `tests/*`	Adds exporter.aimet to fixtures, adds quantize subtest, disables AIMET in a deterministic test, and adjusts loader tests to use `__getitem__`.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

klemen1999
tersekmatija
conorsim

Poem

🐰 I nibbled code lines, snug and neat,

added AIMET treats for models to eat,
loaders now flex — dicts or plain tensors,
ONNX and callbacks hum quantized measures,
hop, export, and ship: a quantized feat.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 36.36% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'AIMET Integration' accurately summarizes the main change—adding AIMET-based quantization support to the project.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests
Commit unit tests in branch feature/aimet

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

kozlov721 added 30 commits March 22, 2026 19:16

aimet

b26e70f

ptq

ba56cb9

Merge branch 'main' into feature/aimet

4bf1bcf

fix for pickle

6e506ce

ignoring quant modules

0e3673b

ptq support

9cd6826

track

3d8828a

more params

7dca296

aimet config

58565e4

requirements

40518c0

serialization

fa33a8d

more ptq techniques

c051144

small fixes

4923c8e

fix keypoint loss

cb07292

optimizer

9dd61d8

config options

d45abc4

loading config values

22c98a1

cleaned up

eafbcf9

docs

59936ec

test

ea07550

fix tests

ea3f8cd

fixes

2716b1e

fixed config

9bccbc8

updated api

d200891

cleanup

d34ca91

cli fix

1140c36

scheduler step remove epoch

fa5c934

fix qat

5180a5e

remmoved hidden node

87f98b5

device switching

ffa5bda

removed requirement

c06d002

kozlov721 commented May 19, 2026

View reviewed changes

kozlov721 added 2 commits May 20, 2026 10:51

fix tests

0162f88

change requirements

bb7f28c

This comment was marked as resolved.

Sign in to view

kozlov721 added 6 commits May 21, 2026 18:39

fix test

25656ab

fix readme

7bcd04b

moved outside of loop

c309012

loading weights before eval

a1fb63b

safer values

4dc9884

fix docs

70faf84

kozlov721 force-pushed the feature/aimet branch from d872f50 to 70faf84 Compare May 21, 2026 16:57

dtronmans approved these changes May 27, 2026

View reviewed changes

Merge branch 'main' into feature/aimet

8d25cf2

This comment was marked as resolved.

Sign in to view

Merge branch 'main' into feature/aimet

84561fa

kozlov721 force-pushed the feature/aimet branch from de26e46 to 84561fa Compare June 3, 2026 13:35

kozlov721 added 9 commits June 5, 2026 17:38

Merge branch 'main' into feature/aimet

de59607

test cli

ac10931

readme update

330ba79

remove an annoying validator

b604804

fix bug

1026ee8

fix weights loading

2345808

disallow empty dataset for quant

fd252bd

Merge branch 'main' into feature/aimet

89c5814

fix

42370f1

kozlov721 merged commit e0e268d into main Jun 8, 2026
16 of 20 checks passed

kozlov721 deleted the feature/aimet branch June 8, 2026 13:59

coderabbitai Bot mentioned this pull request Jun 8, 2026

Advanced Finetuning #358

Open

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

AIMET Integration#366

AIMET Integration#366
kozlov721 merged 89 commits into
mainfrom
feature/aimet

kozlov721 commented Apr 8, 2026 •

edited by coderabbitai Bot

Loading

Uh oh!

kozlov721 left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot commented May 21, 2026 •

edited

Loading

Review limit reached

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

This comment was marked as resolved.

Uh oh!

This comment was marked as resolved.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

kozlov721 commented Apr 8, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Purpose

Specification

Dependencies & Potential Impact

Deployment Plan

Testing & Validation

Summary by CodeRabbit

Uh oh!

kozlov721 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

coderabbitai Bot commented May 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Review limit reached

Walkthrough

Changes

❌ Failed checks (1 warning)

Uh oh!

This comment was marked as resolved.

Uh oh!

This comment was marked as resolved.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

kozlov721 commented Apr 8, 2026 •

edited by coderabbitai Bot

Loading

coderabbitai Bot commented May 21, 2026 •

edited

Loading